from typing import Dict, Any, Mapping
from learners.strategies.training_strategies import (
    TrainingStrategy,
    StandardTraining,
    MixupTraining,
    ClassAugTraining,
)


class TrainingManager:
    """Manages training strategies and their lifecycle"""

    def __init__(self, strategy_config: Dict[str, Any]):
        self.strategies = self._create_strategies(strategy_config)
        self.strategy_schedule = strategy_config.get("schedule", [])

    def _create_strategies(
        self, config: Dict[str, Any]
    ) -> Mapping[str, TrainingStrategy]:
        """Factory method for creating training strategies"""
        strategies: Dict[str, ClassAugTraining | StandardTraining | MixupTraining] = {
            "standard": StandardTraining(),
            "mixup": MixupTraining(alpha=config.get("mixup_alpha", 0.3)),
            "classAug": ClassAugTraining(
                alpha=config.get("alpha", 20),
                mix_times=config.get("mix_times", 4),
                cls_per=config.get("cls_per", 20),
            ),
        }
        return strategies

    def get_strategy_for_task(self, task_id: int) -> TrainingStrategy:
        """Get appropriate strategy for given task"""
        strategy_name = (
            self.strategy_schedule[task_id]
            if task_id < len(self.strategy_schedule)
            else self.strategy_schedule[-1]
        )
        return self.strategies[strategy_name]
